# !/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
from future_hq import FutureHqData
import bcolz
from pandas import DataFrame
import pandas as pd
import sys
sys.path.append("../../")
from bangth_utils.user_logbook import user_log as logger


def create_dir(path):
    while True:
        if not os.path.exists(path):
            os.makedirs(path)


class BcolzSync:
    """
    bcolz操作类
    """
    def __init__(self):
        pass

    @classmethod
    def save_mb(cls, path=None, interval: int = 5):
        default_data_path = os.path.abspath(os.path.expanduser('~/.rqalpha'))
        if path is None:
            path = default_data_path
        else:
            path = os.path.abspath(os.path.join(path, './'))

        mb_bundle = os.path.abspath(os.path.join(path, 'mb_bundle'))
        logger.info(f"mb_bundle: {mb_bundle}")
        if not os.path.exists(mb_bundle):
            os.makedirs(mb_bundle)

        codes = FutureHqData.get_future_codes(interval)

        # hqitem, fields = FutureHqData.get_future_hq(interval)
        mb_index_list = list()
        line_map = {}
        end_date = {}
        index = 0
        s, e = 0, 0
        if interval == "1":
            interval = ""

        dst_path = os.path.abspath(os.path.join(mb_bundle, 'futures_{}mb.bcolz'.format(interval)))
        base_index = 0
        for code in codes:
            pos = code.find(".")
            code = code[pos + 1:] if pos >= 0 else code
            print(code)
            date_dict = {}
            result = FutureHqData.get_future_hq(interval, code)
            for date in result['date']:
                if date not in date_dict:
                    date_dict[date] = 1
                else:
                    date_dict[date] += 1
            # 计算一天的数据索引到mb_index_list中,类似这样的数据
            base_index = cls.calc_mb_indexs(date_dict, mb_index_list, base_index)  # [[20200323, 0, 48], [20200324, 48, 96], [...]]

            e += len(date_dict)
            line_map[code] = [s, e]
            s = e
            end_date[code] = int(list(date_dict.keys())[-1])  # 某支票的最后一天给出来

            arr = DataFrame(result, columns=list(result.dtype.names))
            if index == 0:
                index = 1
                ct = bcolz.ctable.fromdataframe(arr, rootdir=dst_path, mode='w')
                ct.flush()
                ct = bcolz.open(rootdir=dst_path, mode='a')
            else:
                tt = bcolz.ctable.fromdataframe(arr)
                ct.append(tt)

        ct.attrs['end_date'] = end_date
        ct.flush()

        print("保存到" + dst_path + "完成")
        index_df = pd.DataFrame(data=mb_index_list, columns=['date', 'start_at', 'end_at'])
        dst_index_path = "{}/futures_{}mb_index.bcolz".format(mb_bundle, interval)
        index_ct = bcolz.ctable.fromdataframe(index_df, rootdir=dst_index_path, mode='w')
        index_ct.attrs['line_map'] = line_map
        index_ct.flush()
        print("保存到" + dst_index_path + "完成")

    @classmethod
    def calc_mb_indexs(cls, date_dict, mb_index_list, base_index):
        s, e = base_index, base_index
        for date in date_dict.keys():
            e += date_dict[date]
            mb_index_list.append([date, s, e])
            s = e
        return e


if __name__ == '__main__':
    BcolzSync.save_mb()
